/*
 * @Author: Wangjun
 * @Date: 2021-05-15 22:43:20
 * @LastEditTime: 2024-11-20 09:26:49
 * @LastEditors: wangjun haodreams@163.com
 * @Description:
 * @FilePath: \fan_wind_powerd:\go\src\gitee.com\haodreams\golib\wrap\wrap.go
 * hnxr
 */
package wrap

import (
	"sync"
)

// WaitGroupWrapper .
type WaitGroupWrapper struct {
	sync.WaitGroup
	max chan int
}

// NewWaitGroupWrapper 支持设置最大线程数的wg
func NewWaitGroupWrapper(maxThread int) *WaitGroupWrapper {
	wg := new(WaitGroupWrapper)
	if maxThread <= 0 {
		maxThread = 2
	}
	wg.max = make(chan int, maxThread)

	return wg
}

// Wrap cb 回调函数
// 回调函数需要注意的是 回调函数内部的变量是引用变量，需要注意创建新变量复制
func (m *WaitGroupWrapper) Wrap(cb func()) {
	m.max <- len(m.max)
	m.Add(1)
	go func() {
		cb()
		m.Done()
		<-m.max
	}()
}
